﻿using System;
using System.Collections.Generic;
using System.ServiceModel;
using System.ServiceModel.Channels;
using System.ServiceModel.Dispatcher;

namespace SimpleServiceBus.EndpointManagement.Server.Wcf
{
    public class UnitOfWorkCallContextInitializer : ICallContextInitializer
    {
        private readonly Func<IDisposable> _unitOfWorkProvider;
        private readonly Dictionary<Guid, IDisposable> _work = new Dictionary<Guid, IDisposable>();
        private readonly object _threadLock = new object();

        public UnitOfWorkCallContextInitializer(Func<IDisposable> unitOfWorkProvider)
        {
            _unitOfWorkProvider = unitOfWorkProvider;    
        }

        public object BeforeInvoke(InstanceContext instanceContext, IClientChannel channel, Message message)
        {
            if (_unitOfWorkProvider != null)
            {
                IDisposable work = _unitOfWorkProvider();
                if (work != null)
                {
                    Guid correlationObject = Guid.NewGuid();

                    lock(_threadLock)
                        _work.Add(correlationObject, work);


                    return correlationObject;
                }
            }
            return null;
        }

        public void AfterInvoke(object correlationState)
        {
            if (correlationState != null && correlationState is Guid)
            {
                var correlationObject = (Guid) correlationState;

                IDisposable work = null;
                lock (_threadLock)
                {
                    if (_work.ContainsKey(correlationObject))
                    {
                        work = _work[correlationObject];
                        _work.Remove(correlationObject);
                    }
                }


                if (work != null)
                    work.Dispose();
            }
        }
    }
}
